{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Quick Start"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `ssms` package serves two purposes. \n",
    "\n",
    "1. Easy access to *fast simulators of sequential sampling models*\n",
    "   \n",
    "2. Support infrastructure to construct training data for various approaches to likelihood / posterior amortization\n",
    "\n",
    "We provide two minimal examples here to illustrate how to use each of the two capabilities.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Install \n",
    "\n",
    "Let's start with *installing* the `ssms` package.\n",
    "\n",
    "You can do so by typing,\n",
    "\n",
    "`pip install git+https://github.com/AlexanderFengler/ssm_simulators`\n",
    "\n",
    "in your terminal.\n",
    "\n",
    "Below you find a basic tutorial on how to use the package."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Tutorial"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import necessary packages\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import ssms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Using the Simulators\n",
    "\n",
    "Let's start with using the basic simulators. \n",
    "You access the main simulators through the  `ssms.basic_simulators.simulator` function.\n",
    "\n",
    "To get an idea about the models included in `ssms`, use the `config` module.\n",
    "The central dictionary with metadata about included models sits in `ssms.config.model_config`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['ddm',\n",
       " 'ddm_legacy',\n",
       " 'ddm_deadline',\n",
       " 'angle',\n",
       " 'weibull',\n",
       " 'levy',\n",
       " 'levy_angle',\n",
       " 'full_ddm',\n",
       " 'gamma_drift',\n",
       " 'gamma_drift_angle']"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Check included models\n",
    "list(ssms.config.model_config.keys())[:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'name': 'ddm',\n",
       " 'params': ['v', 'a', 'z', 't'],\n",
       " 'param_bounds': [[-3.0, 0.3, 0.1, 0.0], [3.0, 2.5, 0.9, 2.0]],\n",
       " 'boundary': <function ssms.basic_simulators.boundary_functions.constant(t=0)>,\n",
       " 'n_params': 4,\n",
       " 'default_params': [0.0, 1.0, 0.5, 0.001],\n",
       " 'hddm_include': ['z'],\n",
       " 'nchoices': 2}"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Take an example config for a given model\n",
    "ssms.config.model_config[\"ddm\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Note:**\n",
    "The usual structure of these models includes,\n",
    "\n",
    "- Parameter names (`'params'`)\n",
    "- Bounds on the parameters (`'param_bounds'`)\n",
    "- A function that defines a boundary for the respective model (`'boundary'`)\n",
    "- The number of parameters (`'n_params'`)\n",
    "- Defaults for the parameters (`'default_params'`)\n",
    "- The number of choices the process can produce (`'nchoices'`)\n",
    "\n",
    "The `'hddm_include'` key concerns information useful for integration with the [hddm](https://github.com/hddm-devs/hddm) python package, which facilitates hierarchical bayesian inference for sequential sampling models. It is not important for the present tutorial."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ssms.basic_simulators.simulator import simulator\n",
    "\n",
    "sim_out = simulator(\n",
    "    model=\"ddm\", theta={\"v\": 0, \"a\": 1, \"z\": 0.5, \"t\": 0.5}, n_samples=1000\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The output of the simulator is a `dictionary` with three elements.\n",
    "\n",
    "1. `rts` (array)\n",
    "2. `choices` (array)\n",
    "3. `metadata` (dictionary)\n",
    "\n",
    "The `metadata` includes the named parameters, simulator settings, and more."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Using the Training Data Generators\n",
    "\n",
    "The training data generators sit on top of the simulator function to turn raw simulations into usable training data for training machine learning algorithms aimed at posterior or likelihood armortization.\n",
    "\n",
    "We will use the `data_generator` class from `ssms.dataset_generators`. Initializing the `data_generator` boils down to supplying two configuration dictionaries.\n",
    "\n",
    "1. The `generator_config`, concerns choices as to what kind of training data one wants to generate.\n",
    "2. The `model_config` concerns choices with respect to the underlying generative *sequential sampling model*. \n",
    "\n",
    "We will consider a basic example here, concerning data generation to prepare for training [LANs](https://elifesciences.org/articles/65074).\n",
    "\n",
    "Let's start by peeking at an example `generator_config`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'output_folder': 'data/lan_mlp/',\n",
       " 'dgp_list': 'ddm',\n",
       " 'nbins': 0,\n",
       " 'n_samples': 100000,\n",
       " 'n_parameter_sets': 10000,\n",
       " 'n_parameter_sets_rejected': 100,\n",
       " 'n_training_samples_by_parameter_set': 1000,\n",
       " 'max_t': 20.0,\n",
       " 'delta_t': 0.001,\n",
       " 'pickleprotocol': 4,\n",
       " 'n_cpus': 'all',\n",
       " 'kde_data_mixture_probabilities': [0.8, 0.1, 0.1],\n",
       " 'simulation_filters': {'mode': 20,\n",
       "  'choice_cnt': 0,\n",
       "  'mean_rt': 17,\n",
       "  'std': 0,\n",
       "  'mode_cnt_rel': 0.9},\n",
       " 'negative_rt_cutoff': -66.77497,\n",
       " 'n_subruns': 10,\n",
       " 'bin_pointwise': False,\n",
       " 'separate_response_channels': False}"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ssms.config.data_generator_config[\"lan\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You usually have to make just few changes to this basic configuration dictionary.\n",
    "An example below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from copy import deepcopy\n",
    "\n",
    "# Initialize the generator config (for MLP LANs)\n",
    "generator_config = deepcopy(ssms.config.data_generator_config[\"lan\"])\n",
    "# Specify generative model (one from the list of included models mentioned above)\n",
    "generator_config[\"dgp_list\"] = \"angle\"\n",
    "# Specify number of parameter sets to simulate\n",
    "generator_config[\"n_parameter_sets\"] = 100\n",
    "# Specify how many samples a simulation run should entail\n",
    "generator_config[\"n_samples\"] = 1000"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's define our corresponding `model_config`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'name': 'angle', 'params': ['v', 'a', 'z', 't', 'theta'], 'param_bounds': [[-3.0, 0.3, 0.1, 0.001, -0.1], [3.0, 3.0, 0.9, 2.0, 1.3]], 'boundary': <function angle at 0x125e060c0>, 'n_params': 5, 'default_params': [0.0, 1.0, 0.5, 0.001, 0.0], 'hddm_include': ['z', 'theta'], 'nchoices': 2}\n"
     ]
    }
   ],
   "source": [
    "model_config = ssms.config.model_config[\"angle\"]\n",
    "print(model_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are now ready to initialize a `data_generator`, after which we can generate training data using the `generate_data_training_uniform` function, which will use the hypercube defined by our parameter bounds from the `model_config` to uniformly generate parameter sets and corresponding simulated datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "n_cpus used:  6\n",
      "checking:  data/lan_mlp/\n"
     ]
    }
   ],
   "source": [
    "my_dataset_generator = ssms.dataset_generators.lan_mlp.data_generator(\n",
    "    generator_config=generator_config, model_config=model_config\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "simulation round: 1  of 10\n",
      "simulation round: 2  of 10\n",
      "simulation round: 3  of 10\n",
      "simulation round: 4  of 10\n",
      "simulation round: 5  of 10\n",
      "simulation round: 6  of 10\n",
      "simulation round: 7  of 10\n",
      "simulation round: 8  of 10\n",
      "simulation round: 9  of 10\n",
      "simulation round: 10  of 10\n"
     ]
    }
   ],
   "source": [
    "training_data = my_dataset_generator.generate_data_training_uniform(save=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`training_data` is a dictionary containing four keys:\n",
    "\n",
    "1. `data` the features for [LANs](https://elifesciences.org/articles/65074), containing vectors of *model parameters*, as well as *rts* and *choices*.\n",
    "2. `labels` which contain approximate likelihood values\n",
    "3. `generator_config`, as defined above\n",
    "4. `model_config`, as defined above"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can now use this training data for your purposes. If you want to train [LANs](https://elifesciences.org/articles/65074) yourself, you might find the [LANfactory](https://github.com/AlexanderFengler/LANfactory) package helpful.\n",
    "\n",
    "You may also simply find the basic simulators provided with the **ssms** package useful, without any desire to use the outputs into training data for amortization purposes.\n",
    "\n",
    "##### END"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "c2404e761a8d4e2a34f63613cf4c9a9997cd3109cabb959a7904b2035989131a"
  },
  "kernelspec": {
   "display_name": "hddmnn_tutorial",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.4"
  },
  "orig_nbformat": 2
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
